import os
from cv2 import exp
try:
    from .cvio import cvio
except:
    from cvio import cvio
import numpy as np
from copy import deepcopy
import pdb
import shutil

def bbox_iou_matrix(bboxes1, bboxes2, mode='minimum'):
    iou_matrix = np.zeros((len(bboxes2), len(bboxes1)))
    for i in range(len(bboxes2)):
        bboxes = bboxes2[i][None, :]
        bboxes = np.repeat(bboxes, len(bboxes1), axis=0)
        x1 = np.concatenate((bboxes[:, 0:1], bboxes1[:, 0:1]), axis=1).max(1)
        y1 = np.concatenate((bboxes[:, 1:2], bboxes1[:, 1:2]), axis=1).max(1)
        x2 = np.concatenate((bboxes[:, 2:3], bboxes1[:, 2:3]), axis=1).min(1)
        y2 = np.concatenate((bboxes[:, 3:4], bboxes1[:, 3:4]), axis=1).min(1)
        dw = x2 - x1
        dw[dw < 0] = 0
        dh = y2 - y1
        dh[dh < 0] = 0
        inter = dw * dh
        areas1 = (bboxes[:, 2] - bboxes[:, 0]) * (bboxes[:, 3] - bboxes[:, 1])
        areas2 = (bboxes1[:, 2] - bboxes1[:, 0]) * (bboxes1[:, 3] - bboxes1[:, 1])
        if mode == 'standard': 
            union = areas1 + areas2 - inter
        elif mode == 'minimum':
            union = np.concatenate((areas1[:, None], areas2[:, None]), 1).min(1)
        else:
            raise TypeError('Expected mode in ("standard", "minimum"), but got "%s"' % mode)
        iou = inter / union
        iou_matrix[i] = iou
    return iou_matrix.T


def ply2box(points):
    points = np.array(points, dtype=np.float32)
    x1 = points[:, 0].min()
    x2 = points[:, 0].max()
    y1 = points[:, 1].min()
    y2 = points[:, 1].max()
    bbox = [x1, y1, x2, y2]
    return bbox


def load_ann_dict(src1, src2, ann_type='json'):
    print('load annotations(%s) from \n原数据集 %s\n新数据集 %s' % (ann_type, src1, src2))
    annlist = cvio.load_ext_list(src1,ext_type=ann_type,recursive=False)
    annlist = annlist + cvio.load_ext_list(src2, ext_type=ann_type, recursive=False)
    ann_dic = {}
    for ann in annlist:
        name = os.path.basename(ann)
        if not name in ann_dic:
            ann_dic[name] = [ann]
        else:
            ann_dic[name].append(ann)
    print('image files [%d], annotation files [%d]' % (len(ann_dic), len(annlist)))
    return ann_dic


def mergeShapeWithoutOverlap(src1, src2, dst, iouThr=0.8, mode='standard'):
    ann_dic = load_ann_dict(src1, src2, ann_type='json')
    if not os.path.exists(dst):
        os.makedirs(dst)

    for k, (annbase, annlist) in enumerate(ann_dic.items(), 1):
        try:
            output_path = os.path.join(dst, os.path.basename(annbase))
            if len(annlist) == 1:
                shutil.copy(annlist[0], output_path)
                continue
            item_ann = None
            # first step, merge all shapes
            bboxes = []
            for ann in annlist:
                ann_info = cvio.load_ann(ann)
                shapes = deepcopy(ann_info['shapes'])
                bboxes.extend([ply2box(s['points']) for s in shapes])
                if item_ann is None:
                    item_ann = deepcopy(ann_info)
                    continue
                item_ann['shapes'].extend(shapes)
            # second step, filter shapes which are overlaps to anyone other
            bboxes = np.array(bboxes)
            iouM = bbox_iou_matrix(bboxes, bboxes, mode)
            inds2del = []
            for i, iou in enumerate(iouM):
                if i in inds2del: continue
                indsOvlp = list(np.where(iou > iouThr)[0])
                del indsOvlp[indsOvlp.index(i)]
                inds2del.extend(indsOvlp)
            print('Item [%d/%d]: Num origin [%d] Num to delete [%d]' % (k, len(ann_dic), len(bboxes), len(inds2del)))
            item_ann['shapes'] = [s for i,s in enumerate(item_ann['shapes']) if not i in inds2del]
            cvio.write_ann(item_ann,output_path)
        except:
            print('Exception,', annbase)


if __name__ == '__main__':
    # 给入总文件夹 
    # 每个租户的json标注数据放一个子文件夹
    src1 = r'G:\data\drink\hn\images\second_selected\o1' 
    src2 = r'G:\data\drink\hn\images\second_selected\b' 
    # 合并标注存放路径
    dst = r'G:\data\drink\hn\images\second_selected\n'
    mergeShapeWithoutOverlap(src1, src2, dst)
